from flask import Flask, request, jsonify, send_file, make_response
from flask_cors import CORS
import io
import os
import cv2
import torch
import numpy as np
from PIL import Image
import logging
import base64
import sys
sys.path.append('..')
from models import ResnetGenerator
from utils import Preprocess

# 配置日志
logging.basicConfig(level=logging.INFO,
                   format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

class Photo2Cartoon:
    def __init__(self):
        self.pre = Preprocess()
        self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
        self.net = ResnetGenerator(ngf=32, img_size=256, light=True).to(self.device)

        # 加载模型
        params = torch.load('../models/photo2cartoon_weights.pt',
                          map_location=self.device,
                          weights_only=True)
        self.net.load_state_dict(params['genA2B'])
        self.net.eval()

    def inference(self, img):
        # 人脸检测和处理
        face_rgba = self.pre.process(img)
        if face_rgba is None:
            logger.warning('未检测到人脸!')
            return None

        # 调整大小和预处理
        face_rgba = cv2.resize(face_rgba, (256, 256), interpolation=cv2.INTER_AREA)
        face = face_rgba[:, :, :3].copy()
        mask = face_rgba[:, :, 3][:, :, np.newaxis].copy() / 255.
        face = (face * mask + (1 - mask) * 255) / 127.5 - 1

        # 转换为PyTorch张量
        face = np.transpose(face[np.newaxis, :, :, :], (0, 3, 1, 2)).astype(np.float32)
        face = torch.from_numpy(face).to(self.device)

        # 推理
        with torch.no_grad():
            cartoon = self.net(face)[0][0]

        # 后处理
        cartoon = np.transpose(cartoon.cpu().numpy(), (1, 2, 0))
        cartoon = (cartoon + 1) * 127.5
        cartoon = (cartoon * mask + 255 * (1 - mask)).astype(np.uint8)
        return cv2.cvtColor(cartoon, cv2.COLOR_RGB2BGR)

app = Flask(__name__)

# 配置CORS
CORS(app, resources={
    r"/*": {
        "origins": "*",
        "methods": ["GET", "POST", "OPTIONS"],
        "allow_headers": ["Content-Type", "Authorization", "Accept", "Origin"]
    }
})

# 配置上传限制
app.config['MAX_CONTENT_LENGTH'] = 16 * 1024 * 1024  # 16MB 限制

# 创建转换器实例
converter = Photo2Cartoon()

def validate_image(file):
    """验证上传的图片文件"""
    allowed_extensions = {'.png', '.jpg', '.jpeg', '.webp'}
    ext = os.path.splitext(file.filename)[1].lower()
    if ext not in allowed_extensions:
        raise ValueError(f"不支持的文件类型: {ext}")
    
    image_bytes = file.read()
    file.seek(0)
    
    if len(image_bytes) > 10 * 1024 * 1024:  # 10MB
        raise ValueError("文件太大")
        
    try:
        image = Image.open(io.BytesIO(image_bytes))
        image.verify()
        return True
    except Exception as e:
        raise ValueError(f"无效的图像文件: {str(e)}")

@app.route('/upload', methods=['POST', 'OPTIONS'])
def upload_file():
    if request.method == 'OPTIONS':
        response = make_response()
        response.headers.add('Access-Control-Allow-Origin', '*')
        response.headers.add('Access-Control-Allow-Headers', 'Content-Type')
        response.headers.add('Access-Control-Allow-Methods', 'POST')
        return response

    try:
        if 'file' not in request.files:
            return jsonify({"error": "没有上传文件"}), 400
        file = request.files['file']
        if file.filename == '':
            return jsonify({"error": "未选择文件"}), 400
            
        try:
            validate_image(file)
        except ValueError as e:
            return jsonify({"error": str(e)}), 400
            
        # 读取图像
        file_bytes = file.read()
        nparr = np.frombuffer(file_bytes, np.uint8)
        img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        
        logger.info("开始处理图片")
        
        # 转换图片
        result = converter.inference(img)
        
        if result is None:
            return jsonify({"error": "未检测到人脸"}), 400
            
        # 转换为bytes
        _, img_encoded = cv2.imencode('.jpg', result)
        img_bytes = img_encoded.tobytes()
        
        logger.info("图片处理完成")

        # 创建响应
        response = make_response(send_file(
            io.BytesIO(img_bytes),
            mimetype='image/jpeg',
            as_attachment=True,
            download_name='cartoon_style.jpg'
        ))
        
        response.headers['Content-Type'] = 'image/jpeg'
        response.headers['Content-Length'] = len(img_bytes)
        response.headers['Cache-Control'] = 'no-cache'
        response.headers['Access-Control-Allow-Origin'] = '*'
        return response

    except Exception as e:
        logger.error(f"处理过程中出错: {str(e)}", exc_info=True)
        return jsonify({"error": str(e)}), 500

@app.route('/api/uploadimg', methods=['POST', 'OPTIONS'])
def upload_file_api():
    try:
        if 'file' not in request.files:
            return jsonify({"success": False, "error": "没有上传文件"}), 400
        file = request.files['file']
        if file.filename == '':
            return jsonify({"success": False, "error": "未选择文件"}), 400
            
        try:
            validate_image(file)
        except ValueError as e:
            return jsonify({"success": False, "error": str(e)}), 400
            
        # 读取图像
        file_bytes = file.read()
        nparr = np.frombuffer(file_bytes, np.uint8)
        img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        
        logger.info("开始处理图片")
        
        # 转换图片
        result = converter.inference(img)
        
        if result is None:
            return jsonify({"success": False, "error": "未检测到人脸"}), 400
        
        # 转换为base64
        _, img_encoded = cv2.imencode('.jpg', result)
        img_bytes = img_encoded.tobytes()
        image_base64 = base64.b64encode(img_bytes).decode('utf-8')
        
        return jsonify({
            "success": True,
            "resultImagePath": f"data:image/jpeg;base64,{image_base64}"
        })

    except Exception as e:
        logger.error(f"处理过程中出错: {str(e)}", exc_info=True)
        return jsonify({
            "success": False,
            "error": str(e)
        }), 500

@app.route('/api/health', methods=['GET'])
def health_check():
    return jsonify({"status": "ok"}), 200

@app.errorhandler(413)
def request_entity_too_large(error):
    return jsonify({"error": "文件大小超过限制 (最大 16MB)"}), 413

if __name__ == '__main__':
    logger.info("Starting server...")
    app.run(debug=True, host='0.0.0.0', port=3002)
